from ppgnss import gnss_utils
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib.patches as patches
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from mpl_toolkits.axes_grid1 import make_axes_locatable

import utils

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']

width, height = gnss_utils.cm2inch(10), gnss_utils.cm2inch(16)

fontsize=10
filename = "span5_comp_va_5min_down.txt"
data = utils.read_data_file(filename)
snr_thred = 2

damp = data["Damp"]
samp = data["Samp"]
yamp = data["Yamp"]
xamp = data["Xamp"]
uamp = data["Uamp"]

maskd = data["Dsnr"] < snr_thred
masks = data["Ssnr"] < snr_thred

masky = data["Ysnr"] < snr_thred
maskx = data["Xsnr"] < snr_thred
masku = data["Usnr"] < snr_thred

fig, axes = plt.subplots(nrows=3, ncols=1, figsize=(width, height), sharex=False, sharey=False)
plt.subplots_adjust(left=0.12, right=0.99, bottom=0.07, top=0.95, hspace=0.4)
axins = axes[0].inset_axes((0.4, 0.3, 0.4, 0.5))

h1, = axes[0].plot(samp, uamp*1000,  "o", markersize=5, markeredgecolor='grey', markerfacecolor='none',  alpha=0.5)
axins.plot(samp, uamp*1000,  "o", markersize=5, markeredgecolor='grey', markerfacecolor='none',  alpha=0.5)

uamp[masku] = np.nan
h2, = axes[0].plot(samp, uamp*1000, '^', markersize=5, markeredgecolor='none', markerfacecolor='k', alpha=1)
axins.plot(samp, uamp*1000, '^', markersize=5, markeredgecolor='none', markerfacecolor='k', alpha=1)
viewll = (0, 3)
vieww, viewh = 0.5, 4
# rect = patches.Rectangle(viewll, vieww, viewh, linewidth=0.5, edgecolor='k', facecolor='none')
# axes[0][0].add_patch(rect)
mark_inset(axes[0], axins, loc1=2, loc2=4, fc="none", ec='k', lw=1, alpha=0.5)

# 调整子坐标系的显示范围
# axins.set_xlim(xlim0, xlim1)
# axins.set_ylim(ylim0, ylim1)
axins.set_xlim((0, 0.5))
axins.set_ylim((3, 7))
axes[0].set_xlabel("各频振幅 (闸内水位, m)")
axes[0].set_ylabel("各频振幅 (U分量, mm)")
axes[0].set_ylim((0, 15))
axes[0].legend([h1, h2], [u"不显著频点", u"显著频点"], ncol=2, loc='upper left', bbox_to_anchor=(0.1, 1.25))
for x, y, txt in zip(damp, uamp*1000, data["Utide"]):
    if "MM" in txt:
        cx = x-0.05
        cy = y
    elif "S2" in txt:
        cx = x-0.1
        cy = y-1
    elif "MO3" in txt:
        cx, cy = x, y+0.2
    elif "S4" in txt:
        cx, cy = x, y
    elif "MSF" in txt:
        cx, cy = x-0.05, y+0.1
    else:
        cx = x-0.1
        cy = y+1
    if txt not in ["MM", "MSF", "MO3"]:
        axes[0].text(cx, cy, txt)
    else:
        axins.text(cx, cy, txt)

axins = axes[1].inset_axes((0.55, 0.3, 0.35, 0.5))
mark_inset(axes[1], axins, loc1=2, loc2=4, fc="none", ec='k', lw=1, alpha=0.5)
axes[1].plot(damp, yamp*1000, "o", markersize=5, markeredgecolor='grey', markerfacecolor='none',  alpha=0.5)
axins.plot(damp, yamp*1000, "o", markersize=5, markeredgecolor='grey', markerfacecolor='none',  alpha=0.5)

yamp[masky] = np.nan
axes[1].plot(damp, yamp*1000, '^', markersize=5, markeredgecolor='none', markerfacecolor='k', alpha=1)
axins.plot(damp, yamp*1000, '^', markersize=5, markeredgecolor='none', markerfacecolor='k', alpha=1)
viewll = (0, 1.5)
vieww, viewh = 0.4, 3.5
axins.set_xlim((0, 0.35))
axins.set_ylim((1.5, 5))
for x, y, txt in zip(damp, yamp*1000, data["Ytide"]):
    if "MSF" in txt:
        cx, cy = x, y+0.5
    elif "2Q1" in txt:
        cx, cy = x+0.03, y-0.5
    elif "NO1" in txt:
        cx, cy = x, y+0.5
    elif "S4" in txt:
        cx, cy = x+0.03, y-0.5
    elif "SK3" in txt:
        cx, cy = x, y
    else:
        cx, cy = x-0.1, y+1
    if txt in ["K1", "S2", "M2"]:
        axes[1].text(cx, cy, txt)
    else:
        axins.text(cx, cy, txt)
axes[1].set_xlabel("各频振幅 (内外位差, m)")
axes[1].set_ylabel("各频振幅 (Y分量, mm)")
axes[1].set_ylim((0, 15))
# axes[1].legend([h1, h2], [u"不显著频点", u"显著频点"], ncol=2)

axins = axes[2].inset_axes((0.6, 0.2, 0.15, 0.7))
mark_inset(axes[2], axins, loc1=2, loc2=4, fc="none", ec='k', lw=1, alpha=0.5)

axes[2].plot(samp, xamp*1000,  "o", markersize=5, markeredgecolor='grey', markerfacecolor='none',  alpha=0.5)
axins.plot(samp, xamp*1000,  "o", markersize=5, markeredgecolor='grey', markerfacecolor='none',  alpha=0.5)

xamp[maskx] = np.nan
axes[2].plot(samp, xamp*1000, '^', markersize=5, markeredgecolor='none', markerfacecolor='k', alpha=1)
axins.plot(samp, xamp*1000, '^', markersize=5, markeredgecolor='none', markerfacecolor='k', alpha=1)
viewll = (0, 1.5)
vieww, viewh = 0.4, 5.5
axins.set_xlim((-0.02, 0.2))
axins.set_ylim((1.5, 7))
axes[2].set_xlabel("各频振幅 (闸内水位, m)")
axes[2].set_ylabel("各频振幅 (X分量, mm)")
axes[2].set_ylim((0, 15))
# axes[2].legend([h1, h2], [u"不显著频点", u"显著频点"], ncol=2)

for x, y, txt in zip(damp, xamp*1000, data["Xtide"]):
    if txt in ["K1", "S2", "M2"]:
        cx, cy = x-0.1, y+1
        axes[2].text(cx, cy, txt)
    else:
        cx, cy = x+0.02, y
        axins.text(cx, cy, txt)

fig_filename = "figures/tide_echo.png"
plt.savefig(fig_filename, dpi=300)
plt.show()
